package com.orchard.springai.controller;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.util.*;

@CrossOrigin(origins = "*")
@RestController
public class ChatController {


    @Autowired
    private  OpenAiChatModel chatModel;

    // 模拟数据库存储会话和消息
    private final ChatMemory chatMemory = new InMemoryChatMemory();



    @GetMapping("/ai/generate")
    public Map<String,Object> generate(@RequestParam(value = "message") String message) {
        String generation = this.chatModel.call(message);
        // 将生成的结果按字符或词语分隔为列表
        List<String> parts = Arrays.asList(generation.split("(?<=\\G.{5})")); // 每5个字符切分一次
        return Map.of("generation", parts);
    }

//    @GetMapping(value = "/ai/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
//    public SseEmitter generateStream(@RequestParam(value = "message") String message) {
//        SseEmitter emitter = new SseEmitter();
//        Flux<String> stream = chatModel.stream(message);
//
//        stream.subscribe(
//                data -> {
//                    try {
//                        // 将每个数据片段包装成 JSON 并发送
//                        emitter.send(SseEmitter.event().data(Map.of("text", data)));
//                    } catch (Exception e) {
//                        emitter.completeWithError(e);
//                    }
//                },
//                error -> emitter.completeWithError(error),
//                () -> {
//                    try {
//                        // 在流结束时发送一个特殊的 JSON 消息
//                        emitter.send(SseEmitter.event().data(Map.of("finished", true)));
//                        emitter.complete();
//                    } catch (Exception e) {
//                        emitter.completeWithError(e);
//                    }
//                }
//        );
//
//        return emitter;
//    }


//    @GetMapping(value = "/ai/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
//    public SseEmitter generateStream(@RequestParam(value = "message") String message,
//                                     @RequestParam(value = "sessionId", required = false) String sessionId) {
//        SseEmitter emitter = new SseEmitter();
//
//        // 创建 MessageChatMemoryAdvisor 来获取会话历史记录
//        ChatMemory chatMemory = new InMemoryChatMemory(); // 假设这是你实现的内存存储
//        var messageChatMemoryAdvisor = Optional.ofNullable(sessionId)
//                .map(sid -> new MessageChatMemoryAdvisor(chatMemory, sid, 10))
//                .orElse(null);
//
//        Flux<String> stream;
//        if (messageChatMemoryAdvisor != null) {
//            // 如果有会话ID，则使用带有顾问的聊天客户端
//            stream = ChatClient.create(chatModel).prompt()
//                    .user(message)
//                    .advisors(messageChatMemoryAdvisor)
//                    .stream()
//                    .content();
//        } else {
//            // 如果没有会话ID，直接流式调用聊天模型
//            stream = chatModel.stream(message);
//        }
//
//        stream.subscribe(
//                data -> {
//                    try {
//                        // 将每个数据片段包装成 JSON 并发送
//                        emitter.send(SseEmitter.event().data(Map.of("text", data)));
//                    } catch (Exception e) {
//                        emitter.completeWithError(e);
//                    }
//                },
//                error -> emitter.completeWithError(error),
//                () -> {
//                    try {
//                        // 在流结束时发送一个特殊的 JSON 消息
//                        emitter.send(SseEmitter.event().data(Map.of("finished", true)));
//                        emitter.complete();
//                    } catch (Exception e) {
//                        emitter.completeWithError(e);
//                    }
//                }
//        );
//
//        return emitter;
//    }


    /**
     * 根据会话id，从数据库中查找历史消息，并将消息作为上下文回答。
     *
     * @param prompt    用户的提问
     * @param sessionId 会话id
     * @return SSE流响应
     */
//    @GetMapping(value = "chat/stream/history", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
//    public Flux<ServerSentEvent<String>> chatStreamWithHistory(@RequestParam String prompt,
//                                                               @RequestParam String sessionId) {
//        // 1. 如果需要存储会话和消息到数据库，自己可以实现ChatMemory接口，
//        //    这里使用InMemoryChatMemory，内存存储。
//        // 2. 传入会话id，MessageChatMemoryAdvisor会根据会话id去查找消息。
//        // 3. 只需要携带最近10条消息
//        var messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory, sessionId, 10);
//        return ChatClient.create(chatModel).prompt()
//                .user(prompt)
//                // MessageChatMemoryAdvisor会在消息发送给大模型之前，从ChatMemory中获取会话的历史消息，
//                // 然后一起发送给大模型。
//                .advisors(messageChatMemoryAdvisor)
//                .stream()
//                .content()
//                .map(chatResponse -> ServerSentEvent.builder(chatResponse)
//                        .event("message")
//                        .build());
//    }

    @GetMapping(value = "/ai/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> chatStreamWithHistory(
            @RequestParam String message, // 对应前端的message参数
            @RequestParam String sessionId) {

        var messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory, sessionId, 10);

        return ChatClient.create(chatModel).prompt()
                .user(message)
                .advisors(messageChatMemoryAdvisor)
                .stream()
                .content()
                .map(chatResponse -> {
                    // 手动构建JSON字符串
                    String jsonResponse = "{\"text\": \"" + escapeJson(chatResponse) + "\", \"finished\": true}";

                    return ServerSentEvent.<String>builder()
                            .event("message")
                            .data(jsonResponse) // 发送JSON字符串
                            .build();
                });
    }

    // 辅助方法：转义JSON字符串中的特殊字符
    private String escapeJson(String input) {
        if (input == null) return "null";
        return input
                .replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\b", "\\b")
                .replace("\f", "\\f")
                .replace("\n", "\\n")
                .replace("\r", "\\r")
                .replace("\t", "\\t");
    }
}